from __future__ import generator_stop
import os
import argparse
import json
import logging
import threading
import queue

from string import Template
from textwrap import wrap
from elasticsearch import Elasticsearch, helpers

from progressbar import *

from gremlin_python.driver import client
from gremlin_python.process.anonymous_traversal import traversal
from gremlin_python.process.graph_traversal import __, constant, index, label
from gremlin_python.process.traversal import T

LOGGER = logging.getLogger(__name__)

class Labels:

    User = 'USER'
    Tip = 'TIP'
    Review = 'REVIEW'
    Checkin = 'CHECKIN'
    Business = 'BUSINESS'
    Friend = 'FRIEND'


class Loader:

    def __init__(self, connection, datadir) -> None:
        super().__init__()

        self._ConnectionStr = connection
        self._DataDir = datadir
        self._Queue = queue.Queue(10)

    def loadJsonFile(self, name, label, handler, transaction_handler, batchSize=1):

        path = os.path.join(self._DataDir, name)

        fi = open(path, 'rb')
        total = 0
        for i, line in enumerate(open(path, 'rb')):
            total += 1

        progressbar_widgets = ['{}({}) :'.format(label, total), Percentage(), ' ', Bar('#'), ' ', Timer(), ' ', ETA(), ' ', Counter()]
        progressbar_bar = ProgressBar(widgets=progressbar_widgets, maxval=total).start()

        fi = open(path, 'rb')
        count = 0
        actions = []
        while True:
            line = fi.readline()
            if not line:
                break
            count += 1
            progressbar_bar.update(count)

            for action in handler(json.loads(line)):
                actions.append(action)
                if count % batchSize == 0:
                    transaction_handler(actions)
                    actions.clear()
        transaction_handler(actions)
        actions.clear()
        fi.close()

    def show(self):
        pass

    def create(self):
        pass

    def drop(self):
        pass

    def load(self):
        pass


class ElasticsearchLoader(Loader):

    def __init__(self, connection, datadir=None) -> None:
        super().__init__(connection, datadir)

        self._Client = Elasticsearch(hosts=[self._ConnectionStr])
    

    def createIndex(self, name, properties):
        body = {
            "mappings": {
                "properties": properties
            }
        }
        self._Client.indices.create(index=name, body=body)

    def dropIndex(self, name):
        self._Client.indices.delete(index=name)

    def createIndexUser(self):
        name = 'yelp-user'
        properties = {
            'name': {
                'type': 'text',
            }
        }
        self.createIndex(name, properties)
    

    def loadYelpUser(self, name='yelp_academic_dataset_user.json'):

        def yelp_user_handler(data):
            pk = data['user_id']
            properties = {
                "name": data["name"],
            }
            return {
                "_index": "yelp-user",
                "_id": pk,
                "_type": "_doc",
                "_source": properties
            }
        def yelp_user_transaction(actions):
            helpers.bulk(self._Client, actions)
            
        self.loadJsonFile(name, Labels.User, yelp_user_handler, yelp_user_transaction, 2000)
        pass

    def show(self):
       o = self._Client.indices.get_mapping(index="yelp-user", pretty=True)
       print(json.dumps(o, indent=2))
    
    def create(self):
        self.createIndexUser()
    
    def drop(self):
        self.dropIndex("yelp-user")

    def load(self):
        self.loadYelpUser()


class JanusgraphLoader(Loader):

    def __init__(self, connection, datadir=None, traversal_source='g') -> None:
        super().__init__(connection, datadir)

        self._Client = client.Client(connection, traversal_source=traversal_source)
    

    def wrapWithTransaction(self, message):
        return '''
        tx = g.newTransaction()
        {}
        tx.commit()
        '''.format(message)
    
    def wrapWithTraversal(self, message):
        return '''
        g = ConfiguredGraphFactory.open("yelp").traversal()
        {}
        '''.format(message)


    def submitGremlinAsync(self, message, bindings=None):
        return self._Client.submitAsync(message, bindings)


    def showDatabase(self):
        statement = '''
        graph = ConfiguredGraphFactory.open("yelp") \n
        mgmt = graph.openManagement() \n
        mgmt.printSchema() \n
        '''
        return self._Client.submit(statement)

    def createDatabase(self):
        statement = '''
        ConfiguredGraphFactory.create("yelp") \n
        graph = ConfiguredGraphFactory.open("yelp") \n
        mgmt = graph.openManagement() \n
        mgmt.makeEdgeLabel("FRIEND").make() \n
        mgmt.makeVertexLabel("USER").make() \n
        propertyId = mgmt.makePropertyKey('_id_').dataType(String.class).make() \n
        mgmt.makePropertyKey("USER_name").dataType(String.class).make() \n
        mgmt.makePropertyKey("USER_reviewCount").dataType(Integer.class).make() \n
        mgmt.buildIndex("byIdComposite", Vertex.class).addKey(propertyId).buildCompositeIndex() \n
        mgmt.commit()
        '''
        return self._Client.submit(statement)
    
    def dropDatabase(self):
        statement = '''
        ConfiguredGraphFactory.drop("yelp")
        '''
        return self._Client.submit(statement)
    
    def buildPropertyExpr(self, properties):

        exprs = []

        template = Template('''.element().property('${name}', ${value})''')

        for name, v in properties.items():
            value = v
            if isinstance(v, str):
                value = '"{}"'.format(v.replace("'", "\\'").replace('"', '\\"'))
            exprs.append(template.substitute({
                'name': name,
                'value': value
            }))
        
        return ''.join(exprs)


    def addVertexAction(self, label, pk, properties={}):

        template = Template('''
        tx.addVertex('${label}').property('_id_', '${pk}')${properties}
        ''')

        statement = template.substitute({
            "label": label,
            "pk": pk,
            "properties": self.buildPropertyExpr(properties)
        })


        return statement

    
    def addEdgeAction(self, label, fromV, toV, properties={}):
        template = Template('''
        g.addE('${label}')${propertyExpr}.from(g.V().has('_id_', '${fromV}')).to(g.V().has('_id_', '${toV}'))
        ''')

        statement = template.substitute({
            'label': label,
            'fromV': fromV,
            'toV': toV,
            'propertyExpr': self.buildPropertyExpr(properties)
        })
        return statement


    def loadYelpUser(self, name='yelp_academic_dataset_user.json'):


        def yelp_user_vertex_handler(data):
            pk = data['user_id']
            properties = {
                "USER_name": data["name"],
                "USER_reviewCount": data['review_count']
            }
            yield self.addVertexAction(Labels.User, pk, properties)
        
        def yel_user_vertex_transaction(actions):
            self._Queue.put(self.wrapWithTransaction(''.join(actions)))
            # self.submitGremlinAsync(self.wrapWithTransaction(''.join(actions)))

        
        def yelp_user_edge_handler(data):
            pk = data['user_id']
            friends = data['friends']
            for friend in friends.split(','):
               yield self.addEdgeAction(Labels.Friend, pk, friend.strip())
        
        def yel_user_edge_transaction(actions):
            self._Queue.put(self.wrapWithTraversal(''.join(actions)),)
            # self.submitGremlinAsync(self.wrapWithTraversal(''.join(actions)))

        self.loadJsonFile(name, Labels.User, yelp_user_vertex_handler, yel_user_vertex_transaction, 100)
        self.loadJsonFile(name, Labels.Friend, yelp_user_edge_handler, yel_user_edge_transaction, 100)
    
    def _threaded(self):
        while True:
            actions = self._Queue.get(True, 25000)
            if not actions:
                break
            self.submitGremlinAsync(''.join(actions))


    def load(self, c=1):
        # force bind traversal source
        self._Client._traversal_source = 'yelp'


        _threads = []
        t = threading.Thread(target=self.loadYelpUser)
        t.start()

        _threads.append(t)

        for i in range(0, c + 1):
            t = threading.Thread(target=self._threaded)
            t.start()
            _threads.append(t)
        
        for thread in _threads:
            thread.join()
        
        # self.loadYelpUser()


    def show(self):
        r = self.showDatabase()
        print(''.join(r.one()))
    
    def create(self):
        r = self.createDatabase()
        r.one()

    def drop(self):
        r = self.dropDatabase()
        r.one()



def janusgraph_command_delegator(args):

    loader = JanusgraphLoader(connection=args.server, datadir=args.data)
    if args.action == 'show':
        loader.show()
    elif args.action == 'create':
        loader.create()
    elif args.action == 'drop':
        loader.drop()
    elif args.action == 'load':
        loader.load(args.concurrency)


def elasticsearch_command_delegator(args):
    loader = ElasticsearchLoader(connection=args.server, datadir=args.data)
    if args.action == 'show':
        loader.show()
    elif args.action == 'create':
        loader.create()
    elif args.action == 'drop':
        loader.drop()
    elif args.action == 'load':
        loader.load()



if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    subparsers = parser.add_subparsers(title='Yelp Tools',)
    janusgraph_parser = subparsers.add_parser('janusgraph')
    janusgraph_parser.set_defaults(func=janusgraph_command_delegator)
    janusgraph_parser.add_argument('-s', '--server', help='Janusgraph Server', default='ws://127.0.0.1:8182/gremlin')
    janusgraph_parser.add_argument('--data', help='Yelp Data Directory')
    janusgraph_parser.add_argument('--action', help='Database Commands[create, drop, show, load]', default='show')
    janusgraph_parser.add_argument('-c','--concurrency', help='Concurrency', default=1)

    elasticsearch_parser = subparsers.add_parser('elasticsearch')
    elasticsearch_parser.set_defaults(func=elasticsearch_command_delegator)
    elasticsearch_parser.add_argument('-s', '--server', help='Elasticsearch Server', default='http://127.0.0.1:9200')
    elasticsearch_parser.add_argument('--data', help='Yelp Data Directory')
    elasticsearch_parser.add_argument('--action', help='Indexes Commands[create, drop, show, load]', default='show')

    args = parser.parse_args()
    args.func(args)
